from time import clock_settime
from lerobot.policies.diffusion.modeling_diffusion import DiffusionPolicy
from kuavo_train.utils.augmenter import (Augmenter,
                                        DeterministicAugmenterColor,
                                        DeterministicAugmenterGeo4Rgbds,
                                        crop_image,
                                        resize_image)
from torch import Tensor, nn
import torch
from collections import deque
from lerobot.constants import ACTION, OBS_ENV_STATE, OBS_IMAGES, OBS_STATE
from kuavo_train.wrapper.policy.diffusion.DiffusionConfigWrapper import CustomDiffusionConfigWrapper
from lerobot.policies.utils import (
    get_device_from_parameters,
    get_dtype_from_parameters,
    get_output_shape,
    populate_queues,
)

from kuavo_train.wrapper.policy.diffusion.DiffusionModelWrapper import CustomDiffusionModelWrapper
import os, builtins
from pathlib import Path
from typing import TypeVar
from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
from huggingface_hub.errors import HfHubHTTPError
import torchvision
import torchvision.transforms.functional

T = TypeVar("T", bound="CustomDiffusionPolicyWrapper")
OBS_DEPTH = "observation.depth"

class CustomDiffusionPolicyWrapper(DiffusionPolicy):
    def __init__(self,         
                 config: CustomDiffusionConfigWrapper,
                 dataset_stats: dict[str, dict[str, Tensor]] | None = None,
    ):
        super().__init__(config, dataset_stats)
        self.diffusion = CustomDiffusionModelWrapper(config)


    def reset(self):
        """Clear observation and action queues. Should be called on `env.reset()`"""
        self._queues = {
            "observation.state": deque(maxlen=self.config.n_obs_steps),
            "action": deque(maxlen=self.config.n_action_steps),
        }
        if self.config.image_features:
            self._queues["observation.images"] = deque(maxlen=self.config.n_obs_steps)
        if self.config.use_depth and self.config.depth_features:
            self._queues["observation.depth"] = deque(maxlen=self.config.n_obs_steps)
        if self.config.env_state_feature:
            self._queues["observation.environment_state"] = deque(maxlen=self.config.n_obs_steps)

    def select_action(self, batch: dict[str, Tensor]) -> Tensor:
        """Select a single action given environment observations.

        This method handles caching a history of observations and an action trajectory generated by the
        underlying diffusion model. Here's how it works:
          - `n_obs_steps` steps worth of observations are cached (for the first steps, the observation is
            copied `n_obs_steps` times to fill the cache).
          - The diffusion model generates `horizon` steps worth of actions.
          - `n_action_steps` worth of actions are actually kept for execution, starting from the current step.
        Schematically this looks like:
            ----------------------------------------------------------------------------------------------
            (legend: o = n_obs_steps, h = horizon, a = n_action_steps)
            |timestep            | n-o+1 | n-o+2 | ..... | n     | ..... | n+a-1 | n+a   | ..... | n-o+h |
            |observation is used | YES   | YES   | YES   | YES   | NO    | NO    | NO    | NO    | NO    |
            |action is generated | YES   | YES   | YES   | YES   | YES   | YES   | YES   | YES   | YES   |
            |action is used      | NO    | NO    | NO    | YES   | YES   | YES   | NO    | NO    | NO    |
            ----------------------------------------------------------------------------------------------
        Note that this means we require: `n_action_steps <= horizon - n_obs_steps + 1`. Also, note that
        "horizon" may not the best name to describe what the variable actually means, because this period is
        actually measured from the first observation which (if `n_obs_steps` > 1) happened in the past.
        """
        # 与环境交互时无需图像增强
        if ACTION in batch:
            batch.pop(ACTION)

        batch = self.normalize_inputs(batch)

        random_crop = self.config.crop_is_random and self.training
        crop_position = None
        if self.config.image_features:
            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
            for key in self.config.image_features:
                batch[key], crop_position = crop_image(batch[key],target_range=self.config.crop_shape,random_crop=random_crop)
                batch[key] = resize_image(batch[key],target_size=self.config.resize_shape, image_type="rgb")
            batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
        if self.config.use_depth and self.config.depth_features:
            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original

            for key in self.config.depth_features:
                if len(crop_position) == 4:
                    batch[key] = torchvision.transforms.functional.crop(batch[key],*crop_position)
                else:
                    batch[key] = torchvision.transforms.functional.center_crop(batch[key],crop_position)
                batch[key] = resize_image(batch[key],target_size=self.config.resize_shape, image_type="depth")
            batch[OBS_DEPTH] = torch.stack([batch[key] for key in self.config.depth_features], dim=-4)

        # NOTE: It's important that this happens after stacking the images into a single key.
        self._queues = populate_queues(self._queues, batch)

        if len(self._queues[ACTION]) == 0:
            actions = self.predict_action_chunk(batch)
            self._queues[ACTION].extend(actions.transpose(0, 1))

        action = self._queues[ACTION].popleft()
        return action

    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, None]:
        """Run the batch through the model and compute the loss for training or validation."""

        
        random_crop = self.config.crop_is_random and self.training
        crop_position = None
        if self.config.image_features:
            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
            for key in self.config.image_features:
                batch[key], crop_position = crop_image(batch[key],target_range=self.config.crop_shape,random_crop=random_crop)
                batch[key] = resize_image(batch[key],target_size=self.config.resize_shape, image_type="rgb")
            # batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
        if self.config.use_depth and self.config.depth_features:
            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original

            for key in self.config.depth_features:
                if len(crop_position) == 4:
                    batch[key] = torchvision.transforms.functional.crop(batch[key],*crop_position)
                else:
                    batch[key] = torchvision.transforms.functional.center_crop(batch[key],crop_position)
                # print(batch[key].dtype,"~~~~~~~~~~~~~~~~~~")
                batch[key] = resize_image(batch[key],target_size=self.config.resize_shape, image_type="depth")
            # batch[OBS_DEPTH] = torch.stack([batch[key] for key in self.config.depth_features], dim=-4)


        batch = self.normalize_inputs(batch)
        # batch[OBS_DEPTH] = torch.tensor(batch[OBS_DEPTH],dtype=batch[OBS_IMAGES].dtype)
        if self.config.image_features:
            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
            batch[OBS_IMAGES] = torch.stack([batch[key] for key in self.config.image_features], dim=-4)
        if self.config.use_depth and self.config.depth_features:
            batch = dict(batch)  # shallow copy so that adding a key doesn't modify the original
            batch[OBS_DEPTH] = torch.stack([batch[key] for key in self.config.depth_features], dim=-4)
        batch = self.normalize_targets(batch)
        # print(batch[OBS_DEPTH].shape, batch[OBS_DEPTH].max(), batch[OBS_DEPTH].min())
        # print(batch[OBS_IMAGES].shape, batch[OBS_IMAGES].max(), batch[OBS_IMAGES].min())
        # raise ValueError()

        loss = self.diffusion.compute_loss(batch)
        # no output_dict so returning None
        return loss, None
    
    @classmethod
    def from_pretrained(
        cls: builtins.type[T],
        pretrained_name_or_path: str | Path,
        *,
        config: CustomDiffusionConfigWrapper | None = None,
        force_download: bool = False,
        resume_download: bool | None = None,
        proxies: dict | None = None,
        token: str | bool | None = None,
        cache_dir: str | Path | None = None,
        local_files_only: bool = False,
        revision: str | None = None,
        strict: bool = False,
        **kwargs,
    ) -> T:
        """
        The policy is set in evaluation mode by default using `policy.eval()` (dropout modules are
        deactivated). To train it, you should first set it back in training mode with `policy.train()`.
        """
        if config is None:
            config = CustomDiffusionConfigWrapper.from_pretrained(
                pretrained_name_or_path=pretrained_name_or_path,
                force_download=force_download,
                resume_download=resume_download,
                proxies=proxies,
                token=token,
                cache_dir=cache_dir,
                local_files_only=local_files_only,
                revision=revision,
                **kwargs,
            )
        model_id = str(pretrained_name_or_path)
        # print(config)
        instance = cls(config, **kwargs)
        if os.path.isdir(model_id):
            print("Loading weights from local directory")
            model_file = os.path.join(model_id, SAFETENSORS_SINGLE_FILE)
            policy = cls._load_as_safetensor(instance, model_file, config.device, strict)
        else:
            try:
                model_file = hf_hub_download(
                    repo_id=model_id,
                    filename=SAFETENSORS_SINGLE_FILE,
                    revision=revision,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
                    token=token,
                    local_files_only=local_files_only,
                )
                policy = cls._load_as_safetensor(instance, model_file, config.device, strict)
            except HfHubHTTPError as e:
                raise FileNotFoundError(
                    f"{SAFETENSORS_SINGLE_FILE} not found on the HuggingFace Hub in {model_id}"
                ) from e

        policy.to(config.device)
        policy.eval()
        return policy